from Data.rawdata import RawData
from LTE import constant as c
import numpy as np

def SSInd():
    return np.hstack((np.arange(33,64),np.arange(65,65+31)))

def ExtractResources(ind, grid):
    return grid.ravel(order="F")[ind]

# def ofdm_rx_at(buf,offset,nfft,half_cp):
#     cp_len0   = c.LTE_CP_LEN_NORM(1, nfft)
#     rxbuf = np.zeros((nfft,),dtype=np.complex64)
#     rxbuf[:-half_cp] = buf[offset:offset+nfft-half_cp]
#     rxbuf[-half_cp:] = buf[offset-half_cp:offset]
#     return np.fft.fftshift(np.fft.fft(rxbuf))

def ofdm_rx_at(buf,offset,nfft,half_cp):
    cp_len1   = c.LTE_CP_LEN_NORM(1, nfft)
    rxbuf = np.zeros((nfft,),dtype=np.complex64)
    offset += cp_len1
    rxbuf[:-half_cp] = buf[offset:offset+nfft-half_cp]
    rxbuf[-half_cp:] = buf[offset-half_cp:offset]
    return np.fft.fftshift(np.fft.fft(rxbuf))

class LTEInfo:
    def __init__(self) -> None:
        self.pid = None
    def f(self,freq):
        self.freq = freq
        return self
    def cellid(self,id):
        self.cid = id
        return self
    def offset(self,off):
        self.off = off
        return self
    def duplex(self,dd):
        self.dup = dd
        return self
    def set_peak(self,p):
        self.peak = p
        return self
    def port(self,p):
        self.pid = p
        return self
    def __str__(self):
        return f"{self.cid}: {self.dup} {self.peak} @ freq={self.freq} off={self.off} pid={self.pid}"
    
class LTE(RawData):
    def __init__(self,fn,centered=True) -> None:
        RawData.__init__(self,fn)
        self.centered = centered
        self.mu = 0
        if not centered:
            self.dec_factor = int(self.fs/(128*15e3 * 2**self.mu))
        else:
            self.dec_factor = 1
            
    def read(self,start=5e-3,end=30e-3):
        RawData.read(self,start,end)
        if self.centered:
            self.waveform = self.raw

    def power(self,f):
        inx = np.argmin(np.abs(self.fr-f))
        return 10*np.log10(self.fd[inx])
    
    def search(self,start=None,end=None):
        self.fr,self.fd = self.spectrum()
        self.avgd = np.mean(10*np.log10(np.abs(self.fd)))
        if start==None:
            start = self.f_center - 50e6
        if end==None:
            end = self.f_center + 50e6
        sg = int(start/100e3)
        se = int(end/100e3)
        n = se - sg
        self.offset = np.zeros((n,),dtype=np.int64)
        self.nid2 = np.zeros((n,),dtype=np.int32)
        self.peak = np.zeros((n,))
        self.cells = []

        for k in range(n):
            f = start + k*100e3
            self.decimate(f)
            self.waveform /= max(self.waveform.real.max(), self.waveform.imag.max())
            self.sync()
            self.offset[k] = self.timingOffset
            self.peak[k] = self.pss_peak
            self.nid2[k] = self.detected_NID2
            id,slot,peak,ok = self.find_sss(-137)
            if ok:
                off = self.timingOffset+128 if slot==0 else self.timingOffset+9600+128
                self.cells.append(LTEInfo().f(f).cellid(id).duplex("fdd").offset(off).port(0).set_peak(peak))
            id,slot,peak,ok = self.find_sss(-(137*3)-1)
            if ok:
                off = self.timingOffset-137*2-1-960 if slot==0 else self.timingOffset-137*2-1-960*6
                self.cells.append(LTEInfo().f(f).cellid(id).duplex("tdd").offset(off).port(0).set_peak(peak))
        k = np.argmax(self.peak)
        peak = self.peak[k]
        self.pss_all = []

        for k in range(n):
            if self.peak[k]>0.5*peak:
                self.pss_all.append((start+k*100e3,self.nid2[k],self.offset[k],self.peak[k]))
        self.pss_all.sort(key=lambda x: np.abs(x[3]),reverse=True)
        return self.pss_all
    
    def check_fdd(self):
        ret = []
        for freq,nid2,off,peak in self.pss_all:
            self.decimate(freq)
            self.waveform /= max(self.waveform.real.max(), self.waveform.imag.max())
            cid,df,pid,pos,amp = self.afc_fdd(freq)
            # print(cid,df,pid,pos,amp)
            if np.abs(df)<1000:
                ret.append(LTEInfo().f(freq).cellid(cid).duplex("fdd").offset(off).port(pid).set_peak(amp))
        ret.sort(key=lambda x: x.peak, reverse=True)
        self.fdd = ret
        return ret            
            